STAR ATTENTION: EFFICIENT LLM INFERENCE OVER LONG SEQUENCES
作者信息
Shantanu Acharya, Fei Jia, Boris Ginsburg
NVIDIA
链接：[2411.17116] Star Attention: Efficient LLM Inference over Long Sequences
摘要：
Inference with Transformer-based Large Language Models (LLMs) on long sequences is both costly and slow due to the quadratic complexity of the self-attention mechanism. We introduce Star Attention, a two-phase block-sparse approximation that improves computational efficiency by sharding attention across multiple hosts while minimizing communication overhead. In the first phase, the context is processed using blockwise-local attention across hosts, in parallel. In the second phase, query and response tokens attend to all prior cached tokens through sequence-global attention. Star Attention integrates seamlessly with most Transformer-based LLMs trained with global attention, reducing memory requirements and inference time by up to 11x while preserving 95-100% of accuracy.
一句话总结概括
针对长文本的Attention分机器优化
Motivation
创新点或贡献
具体设计
阶段1：Prefill
- Long Context被分为多个连续的块，比如C1，C2，C3
- 每个Block负责的Host都接收到C1的一个副本
- 每个Host生成对应的KV Cache（无跨Host通信）
阶段2：Decode
- 输入一个新的Query，传给该Req的各个Host
- 每个Host在每一层先做局部的Attention计算，然后最后一个Host（文本末端）做跨Host的集合Softmax，生成新的KV Cahce
- 直至最后一层，最后一个Host生成对应的Token
对于所有文本：
对于Context部分，黑色是所关注的区域
实验评估
和Ring Attention的评估
长度越长，加速越明显
精度下降也小于3%，甚至有时候精度更好
和Global Attention对比
基本保持了>95%的准确率，甚至有时候优于baseline。在右下角的图中，作者表示下降如此明显的原因可能是因为生成格式的原因。
Block size对准确率的影响
block size越大，准确率越大，速度越慢。
将seq分为四个block，可能可以在准确率和速度的trade-off中收益最好。按照实验的经验而言，对于超过128K的文本，则将block size固定为32k。对于Llama-3.1-70b这种大模型，则将block限制在16K。
五种BenchMark任务分析
Single-NIAH, Multi-NIAH, Multi-Hop Tracing, Aggregation, and Question Answering (QA).
这里实验都将seq分为4个block。Multi-Hop Tracing性能下降最大，这主要是因为host之间通信的信息不足。Aggregation（比如单词提取等任务）增长很大的原因是因为Star Attention将seq分为局部了，局部摘要能力更强了。
更改anchor block的位置和内容
更改anchor block的长度
anchor block接近block长度，可以获得和global attention相似的结果。
背景
先前工作存在的问题概述
难点
补充背景
思考角度
我如何做这个问题
这个洞见可以引申出其他其他方法吗
迁移到长文本 Auto Scale out的情况